import threading
import time
import typing
from collections import defaultdict
import logging

logger = logging.getLogger(__name__)


class ExpiringSet:
    """
    有时效性的set, 元素超时后就会被移除
    """

    def __init__(self, *, ttl: int = 3, added_callbacks=None, expired_callbacks=None):
        """
        A set with values that expire if they haven't been re-added
        in a given amount of time
        :param ttl: seconds for an item to live
        """
        self._lock = threading.Lock()
        self.ttl = ttl  # 需大于_search的间隔

        self.expire_check = threading.Event()

        # elem --> expire_time
        self._data = {}
        # expire_time --> list[elem]
        self.elems_expire_time = defaultdict(list)
        self.most_recent_expire_time = None  # 最近快过期的过期时间

        threading.Thread(target=self.evict_task, daemon=False).start()

        self.added_callbacks = added_callbacks
        self.expired_callbacks = expired_callbacks

    def evict_task(self):
        logger.info("evict task started")
        while True:
            wait_time = (
                self.most_recent_expire_time - int(time.time())
                if self.most_recent_expire_time
                else None
            )

            if wait_time is not None and wait_time <= 0:
                expired_keys = self.elems_expire_time[self.most_recent_expire_time]
                for key in expired_keys:
                    del self._data[key]
                    logger.debug(f"expired elem '{key}' removed")
                del self.elems_expire_time[self.most_recent_expire_time]
                self.__update_most_recent_expire_time()
                for callback in self.expired_callbacks or []:
                    callback(set(expired_keys))
            else:
                logger.debug(f"wait {wait_time} second to evict expired elements")
                self.expire_check.wait(timeout=wait_time)  # 可以在timeout前返回

            self.expire_check.clear()

    def add(self, elem: any, ttl=None) -> None:
        if ttl == None:
            ttl = self.ttl
        elem_exist = elem in self._data
        with self._lock:
            self.__remove_elem_expire_time(elem)

            expire_time = int(time.time() + ttl)
            self._data[elem] = expire_time
            self.elems_expire_time[expire_time].append(elem)

            self.__update_most_recent_expire_time()
            self.expire_check.set()
            logger.info(f'-------------------- add {elem}')

        if not elem_exist:
            for callback in self.added_callbacks or []:
                callback(set([elem]))


    def remove(self, elem: typing.Any) -> typing.Any:
        with self._lock:
            if elem in self._data:
                self.__remove_elem_expire_time(elem)
                del self._data[elem]
                self.__update_most_recent_expire_time()
                self.expire_check.set()
            return elem

    def clear(self) -> None:
        with self._lock:
            self._data = {}
            self.elems_expire_time = defaultdict(list)
            self.expire_check.clear()

    def __remove_elem_expire_time(self, elem):
        elem_expire_time = self._data.get(elem, None)
        if elem_expire_time is not None:
            elems = self.elems_expire_time[elem_expire_time]
            if elem in elems:
                elems.remove(elem)
                if len(elems) == 0:
                    del self.elems_expire_time[elem_expire_time]

    def __update_most_recent_expire_time(self):
        expire_time_list = sorted(self.elems_expire_time.keys())
        self.most_recent_expire_time = (
            expire_time_list[0] if len(expire_time_list) > 0 else None
        )

    def get_eles(self):
        return {k for k in self._data.keys()}

    def __contains__(self, elem: typing.Any) -> bool:
        with self._lock:
            return elem in self._data

    def __iter__(self):
        with self._lock:
            return iter(self.get_eles())
